from dpdata import LabeledSystem,MultiSystems,System
import glob
import os
from subprocess import Popen,PIPE
import numpy as np

MLFF_path = "MLFF"
DFT_path = "DFT"

MLFF_outcar_path = []
for root, dirs, files in os.walk(MLFF_path, topdown=False):
    for name in files:
        if "OUTCAR" == name:
            MLFF_outcar_path.append((root, name))

DFT_outcar_path = []
for root, dirs, files in os.walk(DFT_path, topdown=False):
    for name in files:
        if "OUTCAR" == name:
            DFT_outcar_path.append((root, name))

MLFF_basename = [os.path.basename(i[0]) for i in MLFF_outcar_path]
DFT_basename = [os.path.basename(i[0]) for i in DFT_outcar_path]
for i in MLFF_basename:
    if i not in DFT_basename:
        raise SystemError("Cannot find MLFF-%s in DFT files!" %(i))

MLFF_ms=MultiSystems()
DFT_ms=MultiSystems()
for index, basename in enumerate(MLFF_basename):
    try:
        p2 = Popen('if [ $(grep NELMIN {0} | wc -l) -eq 2 ];then sed -i "0,/NELMIN/s/NELMIN/TMP/" {0};fi ;sed -i "s/free  energy ML TOTEN/free  energy   TOTEN/g" {0};' .format(os.path.join(MLFF_outcar_path[index][0],MLFF_outcar_path[index][1])) ,shell=True,stdout=PIPE)
        p2.wait()  
        tmp_file = os.path.join(MLFF_outcar_path[index][0],MLFF_outcar_path[index][1])
        command = 'grep "VOLUME and BASIS-vectors are now" %s; if [ "$?" -eq "1" ];then num=$(grep -n "volume of cell :" %s | tail -n 1 | awk \'{split($0,a,":");print a[1]}\') ; sed "$num ienergy-cutoff" -i %s ; sed "$num i----------" -i %s;sed "$num i VOLUME and BASIS-vectors are now :" -i %s;fi' %(tmp_file,tmp_file ,tmp_file, tmp_file ,tmp_file )
        p2 = Popen(command,shell=True,stdout=PIPE)
        p2.wait()      
        MLFF_ls=LabeledSystem(os.path.join(MLFF_outcar_path[index][0],MLFF_outcar_path[index][1]))
        POSCAR = System(os.path.join(MLFF_outcar_path[index][0],'POSCAR'))
        MLFF_ls.data['atom_names']=POSCAR.data['atom_names']
        tmp = DFT_outcar_path[DFT_basename.index(basename)]
        p2 = Popen('if [ $(grep NELMIN {0} | wc -l) -eq 2 ];then sed -i "0,/NELMIN/s/NELMIN/TMP/" {0};fi' .format(os.path.join(tmp[0],tmp[1])) ,shell=True,stdout=PIPE)
        p2.wait()   
        DFT_ls=LabeledSystem(os.path.join(tmp[0],tmp[1]))
        POSCAR = System(os.path.join(tmp[0],'POSCAR'))
        DFT_ls.data['atom_names']=POSCAR.data['atom_names']        
    except:
        print(MLFF_outcar_path[index][0],"has errors.")
    if len(DFT_ls)>0 and len(MLFF_ls) > 0:
        MLFF_ms.append(MLFF_ls)
        DFT_ms.append(DFT_ls)
MLFF_ms.to_deepmd_npy('MLFF_deepmd',set_size=MLFF_ms.get_nframes())
DFT_ms.to_deepmd_npy('DFT_deepmd',set_size=DFT_ms.get_nframes()) 

MLFF_system = LabeledSystem(glob.glob("MLFF_deepmd/*")[0],fmt='deepmd/npy')       
DFT_system = LabeledSystem(glob.glob("DFT_deepmd/*")[0],fmt='deepmd/npy')       
np.savetxt("E.txt",np.array([DFT_system['energies'].flatten(),MLFF_system['energies']]).T)
np.savetxt("f.txt",np.concatenate((DFT_system['forces'].reshape(-1,3),MLFF_system['forces'].reshape(-1,3)),axis=1))
np.savetxt("virial.txt",np.concatenate((DFT_system['virials'].reshape(-1,9),MLFF_system['virials'].reshape(-1,9)),axis=1))
data=np.loadtxt("E.txt")
if len(data.shape) == 1:
    data = data[np.newaxis,:]
rmse_E = np.sqrt(np.mean((data[:,0]-data[:,1])**2))  
data=np.loadtxt("f.txt")
data0=np.sqrt(np.sum(data[:,:3]**2,axis=1))
data1=np.sqrt(np.sum(data[:,3:]**2,axis=1))
if len(data.shape) == 1:
    data0 = data0[np.newaxis,:]
    data1 = data1[np.newaxis,:]
rmse_f = np.sqrt(np.mean((data0-data1)**2))
data=np.loadtxt("virial.txt")
if len(data.shape) == 1:
    data = data[np.newaxis,:]
rmse_virial = np.sqrt(np.mean((data[:,:9]-data[:,9:])**2))
print("Energy RMSE        : %f eV" %(rmse_E))
print("Force  RMSE        : %f eV/A" %(rmse_f))
print("Virial RMSE        : %f eV" %(rmse_virial))
